#include "pdtdw.hpp"

// -----------------------------------------------------------
// PDTD (Pixel Depth Test & Dispatch) 模块实现
// 职责：从HR接收像素，过滤无效像素，将有效的像素推入DDR写队列
// -----------------------------------------------------------
void pixel_depth_test_and_dispatch(
    hls::stream<pixel_fragment_packet_fixed>& input_pixel_fragment_stream,
    hls::stream<ddr_write_packet_fixed>& output_ddr_write_stream,
    const ap_uint<11> max_x_res,
    const ap_uint<11> max_y_res,
    const ap_uint<32> num_fragments_from_hr // 从HR接收的像素片段总数
) {
    #pragma HLS INTERFACE axis port=input_pixel_fragment_stream
    #pragma HLS INTERFACE axis port=output_ddr_write_stream 

    #pragma HLS INTERFACE s_axilite port=max_x_res bundle=ctrl_s_axilite
    #pragma HLS INTERFACE s_axilite port=max_y_res bundle=ctrl_s_axilite
    #pragma HLS INTERFACE s_axilite port=num_fragments_from_hr bundle=ctrl_s_axilite 
    #pragma HLS INTERFACE s_axilite port=return bundle=ctrl_s_axilite

    for (ap_uint<32> i = 0; i < num_fragments_from_hr; ++i) { 
        #pragma HLS PIPELINE II=1
        
        pixel_fragment_packet_fixed fragment;
        input_pixel_fragment_stream.read(fragment); 

        if (fragment.is_valid_pixel) {
            ddr_write_packet_fixed write_op;
            write_op.screen_x = fragment.screen_x;
            write_op.screen_y = fragment.screen_y;
            write_op.new_z     = fragment.interpolated_z;
            write_op.pure_color_R = fragment.pure_color_R;
            write_op.pure_color_G = fragment.pure_color_G;
            write_op.pure_color_B = fragment.pure_color_B;

            output_ddr_write_stream.write(write_op);
        }
    }
}


// -----------------------------------------------------------
// DDR_Write_Engine 模块实现
// 职责：从写队列接收指令，执行DDR的读-修改-写操作
// 通过两个独立的 AXI HP 端口访问深度缓冲区和帧缓冲区
// -----------------------------------------------------------
void ddr_write_engine(
    hls::stream<ddr_write_packet_fixed>& input_ddr_write_stream,
    const ap_uint<11> max_x_res,
    const ap_uint<11> max_y_res,
    const ap_uint<32> num_write_ops_to_process, // 接收到PDTDW输出的有效像素数
    zbuffer_ddr_word_t* output_zbuffer_ddr,  // 指向深度缓冲区的 64 位 AXI Master
    color_ddr_word_t* output_framebuffer_ddr // 指向帧缓冲区的 64 位 AXI Master
) {
    // S_AXILITE 接口，用于从PS接收控制参数
    #pragma HLS INTERFACE s_axilite port=max_x_res bundle=ctrl_s_axilite_dwe
    #pragma HLS INTERFACE s_axilite port=max_y_res bundle=ctrl_s_axilite_dwe
    #pragma HLS INTERFACE s_axilite port=num_write_ops_to_process bundle=ctrl_s_axilite_dwe
    #pragma HLS INTERFACE s_axilite port=return bundle=ctrl_s_axilite_dwe

    // AXI4 Full Master 接口，连接到DDR。使用不同的bundle来确保独立的DDR访问路径。
    // Z buffer uses HP0, Framebuffer uses HP1 as per project description.
    #pragma HLS INTERFACE m_axi port=output_zbuffer_ddr  offset=slave bundle=m_axi_gmem_zb_rw max_read_burst_length=8 max_write_burst_length=8
    #pragma HLS INTERFACE m_axi port=output_framebuffer_ddr offset=slave bundle=m_axi_gmem_fb_write max_read_burst_length=8 max_write_burst_length=8

    #pragma HLS INTERFACE axis port=input_ddr_write_stream

    #pragma HLS DATAFLOW // Enable dataflow to pipeline read/write operations from/to DDR.

    // Constants for 64-bit word access
    const ap_uint<32> Z_PER_WORD = 4; // 64 bits / 16 bits = 4 depth values per word
    const ap_uint<32> COLOR_PER_WORD = 2; // 64 bits / 32 bits = 2 color values per word (assuming 24-bit RGB packed into 32-bit slot)
    
    // Masks for 16-bit depth values within a 64-bit word
    const ap_uint<64> Z_MASK_16BIT = 0xFFFF;

    WRITE_OPS_LOOP:
    for (ap_uint<32> i = 0; i < num_write_ops_to_process; ++i) {
        #pragma HLS PIPELINE II=1 // HLS will attempt II=1, potentially with higher latency

        ddr_write_packet_fixed write_op;
        input_ddr_write_stream.read(write_op); // 从队列读取写操作指令

        ap_uint<32> pixel_linear_idx = write_op.screen_y * max_x_res + write_op.screen_x;

        // --- Z-Buffer Access (Read-Modify-Write) ---
        // Calculate 64-bit word index and 16-bit depth index within the word
        ap_uint<32> z_word_idx = pixel_linear_idx / Z_PER_WORD;
        ap_uint<3> z_sub_word_offset = pixel_linear_idx % Z_PER_WORD; // 0, 1, 2, 3

        // Read the 64-bit word containing the old Z value
        zbuffer_ddr_word_t current_z_word = output_zbuffer_ddr[z_word_idx];

        // Extract the specific 16-bit old_z
        fpal::fp_z_buffer_t old_z = (fpal::fp_z_buffer_t)((current_z_word >> (z_sub_word_offset * 16)) & Z_MASK_16BIT);

        // Depth Test: if new pixel is closer (smaller Z value), update
        if (write_op.new_z < old_z) {
            // Modify the 64-bit word with the new_z
            zbuffer_ddr_word_t new_z_word = current_z_word;
            new_z_word &= ~(Z_MASK_16BIT << (z_sub_word_offset * 16)); // Clear old Z bits
            new_z_word |= ((zbuffer_ddr_word_t)write_op.new_z << (z_sub_word_offset * 16)); // Set new Z bits
            
            // Write back the modified 64-bit word to Z-buffer
            output_zbuffer_ddr[z_word_idx] = new_z_word;

            // --- Framebuffer Access (Read-Modify-Write) ---
            // Each 24-bit RGB pixel will be packed into a 32-bit slot within the 64-bit word
            // So, 2 pixels per 64-bit word.
            ap_uint<32> color_word_idx = pixel_linear_idx / COLOR_PER_WORD;
            ap_uint<1> color_sub_word_offset = pixel_linear_idx % COLOR_PER_WORD; // 0 or 1 for pixel 0 or pixel 1

            // Read the 64-bit word containing the old color value(s)
            color_ddr_word_t current_color_word = output_framebuffer_ddr[color_word_idx];

            // Pack new color (RGB888) into a 32-bit slot (e.g., 0x00RRGGBB)
            ap_uint<32> new_packed_color = 0;
            new_packed_color |= ((ap_uint<32>)write_op.pure_color_R << 16);
            new_packed_color |= ((ap_uint<32>)write_op.pure_color_G << 8);
            new_packed_color |= ((ap_uint<32>)write_op.pure_color_B << 0);

            // Modify the 64-bit word with the new packed color
            color_ddr_word_t modified_color_word = current_color_word;
            // Define 32-bit mask for a pixel slot
            const ap_uint<64> PIXEL_SLOT_MASK = 0xFFFFFFFF; 

            if (color_sub_word_offset == 0) { // First pixel slot (lower 32 bits)
                modified_color_word &= ~(PIXEL_SLOT_MASK << 0);
                modified_color_word |= ((color_ddr_word_t)new_packed_color << 0);
            } else { // Second pixel slot (upper 32 bits)
                modified_color_word &= ~(PIXEL_SLOT_MASK << 32);
                modified_color_word |= ((color_ddr_word_t)new_packed_color << 32);
            }
            
            // Write back the modified 64-bit word to Framebuffer
            output_framebuffer_ddr[color_word_idx] = modified_color_word;
        }
    }
}
